{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run -i ../utilsEnhanced.py\n",
    "%run -i ../Tiramisu_ProgramLocal.py\n",
    "%run -i ../Environment_sparse_enhanced.py\n",
    "%run -i ../ActionEnhanced.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tiramisu_path = '/home/narimane/tiramisu' # Put the path to your tiramisu installation here\n",
    "os.environ['TIRAMISU_ROOT'] = tiramisu_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "in init\n",
      "*************************** /home/narimane/Desktop/PFE/RL_in_Search_space_exploration/Notebooks\n"
     ]
    }
   ],
   "source": [
    "\n",
    "env = SearchSpaceSparseEnhanced(\"../../new_dataset.json\", \"../../Dataset\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "in reset\n",
      "directory removed\n",
      "../../Dataset_copies/function550516/function550516_generator.cpp\n",
      "in repr template\n",
      "directory removed\n",
      "../../Dataset_copies/function550086/function550086_generator.cpp\n",
      "in repr template\n",
      "the length is 784\n",
      "the initial execution time exists\n",
      "The initial execution time is 84.2221\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'representation': array([  0.,  64.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   1., 385.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   1., 321.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   1.,   1.,   1.,   1.,   1.,   0.,   0.,\n",
       "          1.,   1.,   1.,   0.,   0.,   0.,   0.,   0.,   1.,   0.,   1.,\n",
       "          0.,   0.,   0.,   0.,   1.,   0.,   0.,   1.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   1.,   1.,   1.,   1.,   1.,   0.,   0.,   1.,\n",
       "          1.,   1.,   0.,   0.,   0.,   0.,   0.,   1.,   0.,   1.,   0.,\n",
       "          0.,   0.,   0.,   1.,   0.,   0.,   1.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   1.,   1.,   1.,   1.,   1.,   0.,   0.,   1.,   1.,\n",
       "          1.,   0.,   0.,   0.,   0.,   0.,   1.,   0.,   1.,   0.,   0.,\n",
       "          0.,  -1.,   1.,   0.,   0.,   1.,   0.,   0.,  -1.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   1.,   1.,   1.,   1.,   1.,   0.,   0.,   1.,   1.,   1.,\n",
       "          0.,   0.,   0.,   0.,   0.,   1.,   0.,   1.,   0.,   0.,   0.,\n",
       "          0.,   1.,   0.,   0.,   1.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          1.,   1.,   1.,   1.,   1.,   0.,   0.,   1.,   1.,   1.,   0.,\n",
       "          0.,   0.,   0.,   0.,   1.,   0.,   1.,   0.,   0.,   0.,   0.,\n",
       "          1.,   0.,   0.,   1.,   0.,   0.,  -1.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   1.,\n",
       "          1.,   1.,   1.,   1.,   0.,   0.,   1.,   1.,   1.,   0.,   0.,\n",
       "          0.,   0.,   0.,   1.,   0.,   1.,   0.,   0.,   0.,   0.,   1.,\n",
       "          0.,   0.,   1.,   0.,   0.,   1.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   1.,   1.,\n",
       "          1.,   1.,   1.,   0.,   0.,   1.,   1.,   1.,   0.,   0.,   0.,\n",
       "          0.,   0.,   1.,   0.,   1.,   0.,   0.,   0.,  -1.,   1.,   0.,\n",
       "          0.,   1.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   1.,   1.,   1.,\n",
       "          1.,   1.,   0.,   0.,   1.,   1.,   1.,   0.,   0.,   0.,   0.,\n",
       "          0.,   1.,   0.,   1.,   0.,   0.,   0.,  -1.,   1.,   0.,   0.,\n",
       "          1.,   0.,   0.,   1.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   1.,   1.,   1.,   1.,\n",
       "          1.,   0.,   0.,   1.,   1.,   1.,   0.,   0.,   0.,   0.,   0.,\n",
       "          1.,   0.,   1.,   0.,   0.,   0.,   1.,   1.,   0.,   0.,   1.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   1.,   1.,   1.,   1.,   1.,\n",
       "          0.,   0.,   1.,   1.,   1.,   0.,   0.,   0.,   0.,   0.,   1.,\n",
       "          0.,   1.,   0.,   0.,   0.,   1.,   1.,   0.,   0.,   1.,   0.,\n",
       "          0.,  -1.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   1.,   1.,   1.,   1.,   1.,   0.,\n",
       "          0.,   1.,   1.,   1.,   0.,   0.,   0.,   0.,   0.,   1.,   0.,\n",
       "          1.,   0.,   0.,   0.,   1.,   1.,   0.,   0.,   1.,   0.,   0.,\n",
       "          1.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,\n",
       "          0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   6.,\n",
       "          1.,   2.,   0.], dtype=float32),\n",
       " 'action_mask': array([1., 1., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 0., 0., 0., 0.,\n",
       "        0., 1., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 0., 0., 0., 0.,\n",
       "        0., 1.], dtype=float32)}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3\n",
      "{0: {'iterator': 'i0', 'lower_bound': 0, 'upper_bound': 64}, 1: {'iterator': 'i1', 'lower_bound': 1, 'upper_bound': 385}, 2: {'iterator': 'i2', 'lower_bound': 1, 'upper_bound': 321}}\n"
     ]
    }
   ],
   "source": [
    "print(len(env.it_dict))\n",
    "print(env.it_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(env.step(41))\n",
    "print(env.obs['action_mask'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "in step function\n",
      "the curr schedule is:  0 \n",
      "inisde try\n",
      "after creating the action\n",
      "in action params function\n",
      "after choosing first and second params and factors\n",
      "\n",
      " first loop extent is  64\n",
      "first factor is 64\n",
      "tiling flag 1 false, loopextent == factor\n",
      "\n",
      " second loop extent is  384\n",
      "second factor is 64\n",
      "action params first are {'tiling_depth': 2, 'first_dim_index': 0, 'tiling_loop_1': False, 'second_dim_index': 1, 'tiling_loop_2': True, 'first_factor': 64, 'second_factor': 64, 'tiling_loop_3': None, 'third_factor': None}\n",
      "\n",
      " The chosen action is:  28\n",
      "in tiling\n",
      "\n",
      " ******* flags ****** False True\n",
      "\n",
      " *-*-*- Legality Check File -*-*-* \n",
      "\n",
      "#include <tiramisu/tiramisu.h> \n",
      "#include <tiramisu/auto_scheduler/evaluator.h>\n",
      "#include <tiramisu/auto_scheduler/search_method.h>\n",
      "\n",
      "using namespace tiramisu;\n",
      "\n",
      "int main(int argc, char **argv){                \n",
      "\ttiramisu::init(\"function550086\");\n",
      "\tvar i0(\"i0\", 0, 64), i1(\"i1\", 1, 385), i2(\"i2\", 1, 321), i1_p1(\"i1_p1\", 0, 386), i2_p1(\"i2_p1\", 0, 322);\n",
      "\tinput icomp00(\"icomp00\", {i0,i1_p1,i2_p1}, p_float64);\n",
      "\tcomputation comp00(\"comp00\", {i0,i1,i2},  p_float64);\n",
      "\tcomp00.set_expression(icomp00(i0, i1, i2)*icomp00(i0, i1 - 1, i2 - 1) + icomp00(i0, i1, i2) + icomp00(i0, i1, i2 - 1) + icomp00(i0, i1, i2 + 1) - icomp00(i0, i1 - 1, i2) + icomp00(i0, i1 - 1, i2 + 1)*icomp00(i0, i1 + 1, i2) + icomp00(i0, i1 + 1, i2 - 1) + icomp00(i0, i1 + 1, i2 + 1));\n",
      "\tbuffer buf00(\"buf00\", {64,386,322}, p_float64, a_output);\n",
      "\ticomp00.store_in(&buf00);\n",
      "\tcomp00.store_in(&buf00);\n",
      "\t\n",
      "    prepare_schedules_for_legality_checks();\n",
      "    perform_full_dependency_analysis();\n",
      "    \n",
      "    bool is_legal=true;\n",
      "\tcomp00.tile(0,1,64,64);\n",
      "\n",
      "    is_legal &= check_legality_of_function();\n",
      "    \n",
      "    std::ofstream out(\"../../Dataset_copies/function550086/legality_check_result.txt\");\n",
      "    out << is_legal;\n",
      "    out.close();\n",
      "         \n",
      "\treturn 0; \n",
      "}\n",
      "\n",
      " in tiling,  lc res: 1\n",
      "inside loop tiling 2\n",
      "after loop tiling 2\n",
      "\n",
      " tiling applied\n",
      "the original iterators were: {0: {'iterator': 'i0', 'lower_bound': 0, 'upper_bound': 64}, 1: {'iterator': 'i1', 'lower_bound': 1, 'upper_bound': 385}, 2: {'iterator': 'i2', 'lower_bound': 1, 'upper_bound': 321}}\n",
      "tiling depth is  2\n",
      "keys:  [0, 1, 2]\n",
      "first depth is 0\n",
      "second depth is 1\n",
      "len of action_params is 9\n",
      "it_dict after popping only {0: {'iterator': 'i0', 'lower_bound': 0, 'upper_bound': 64}, 1: {'iterator': 'i1', 'lower_bound': 1, 'upper_bound': 385}, 2: {'iterator': 'i2', 'lower_bound': 1, 'upper_bound': 321}, 3: {'iterator': 'i2', 'lower_bound': 1, 'upper_bound': 321}}\n",
      "it_dict after pop only {0: {'iterator': 'i0', 'lower_bound': 0, 'upper_bound': 64}, 1: {'iterator': 'i1', 'lower_bound': 1, 'upper_bound': 385}, 2: {'iterator': 'i2', 'lower_bound': 1, 'upper_bound': 321}, 3: {'iterator': 'i2', 'lower_bound': 1, 'upper_bound': 321}}\n",
      "in action params == 7 and tiling_loop_2\n",
      "after update iterators with the schedule T2(L0,L1,64,64) it is {0: {'iterator': 'i0', 'lower_bound': 0, 'upper_bound': 64}, 1: {'iterator': 'i1', 'lower_bound': 1, 'upper_bound': 6.015625}, 2: {'iterator': 'i1_1', 'lower_bound': 1, 'upper_bound': 64}, 3: {'iterator': 'i2', 'lower_bound': 1, 'upper_bound': 321}}\n",
      "--- done is ---- False\n",
      "the reward is 0\n",
      "({'representation': array([  0.      ,  64.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   1.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,  64.      ,   0.      ,   1.      ,   6.015625,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   1.      ,\n",
      "         0.      ,   0.      ,   0.      ,  64.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,  64.      ,\n",
      "         0.      ,   1.      , 321.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   1.      ,   1.      ,   1.      ,\n",
      "         1.      ,   1.      ,   0.      ,   0.      ,   1.      ,\n",
      "         1.      ,   1.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   1.      ,   0.      ,   1.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   1.      ,\n",
      "         0.      ,   0.      ,   1.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         1.      ,   1.      ,   1.      ,   1.      ,   1.      ,\n",
      "         0.      ,   0.      ,   1.      ,   1.      ,   1.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         1.      ,   0.      ,   1.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   1.      ,   0.      ,   0.      ,\n",
      "         1.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   1.      ,   1.      ,\n",
      "         1.      ,   1.      ,   1.      ,   0.      ,   0.      ,\n",
      "         1.      ,   1.      ,   1.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   1.      ,   0.      ,\n",
      "         1.      ,   0.      ,   0.      ,   0.      ,  -1.      ,\n",
      "         1.      ,   0.      ,   0.      ,   1.      ,   0.      ,\n",
      "         0.      ,  -1.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   1.      ,   1.      ,   1.      ,   1.      ,\n",
      "         1.      ,   0.      ,   0.      ,   1.      ,   1.      ,\n",
      "         1.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   1.      ,   0.      ,   1.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   1.      ,   0.      ,\n",
      "         0.      ,   1.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   1.      ,\n",
      "         1.      ,   1.      ,   1.      ,   1.      ,   0.      ,\n",
      "         0.      ,   1.      ,   1.      ,   1.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   1.      ,\n",
      "         0.      ,   1.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   1.      ,   0.      ,   0.      ,   1.      ,\n",
      "         0.      ,   0.      ,  -1.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   1.      ,   1.      ,   1.      ,\n",
      "         1.      ,   1.      ,   0.      ,   0.      ,   1.      ,\n",
      "         1.      ,   1.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   1.      ,   0.      ,   1.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   1.      ,\n",
      "         0.      ,   0.      ,   1.      ,   0.      ,   0.      ,\n",
      "         1.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         1.      ,   1.      ,   1.      ,   1.      ,   1.      ,\n",
      "         0.      ,   0.      ,   1.      ,   1.      ,   1.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         1.      ,   0.      ,   1.      ,   0.      ,   0.      ,\n",
      "         0.      ,  -1.      ,   1.      ,   0.      ,   0.      ,\n",
      "         1.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   1.      ,   1.      ,\n",
      "         1.      ,   1.      ,   1.      ,   0.      ,   0.      ,\n",
      "         1.      ,   1.      ,   1.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   1.      ,   0.      ,\n",
      "         1.      ,   0.      ,   0.      ,   0.      ,  -1.      ,\n",
      "         1.      ,   0.      ,   0.      ,   1.      ,   0.      ,\n",
      "         0.      ,   1.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   1.      ,   1.      ,   1.      ,   1.      ,\n",
      "         1.      ,   0.      ,   0.      ,   1.      ,   1.      ,\n",
      "         1.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   1.      ,   0.      ,   1.      ,   0.      ,\n",
      "         0.      ,   0.      ,   1.      ,   1.      ,   0.      ,\n",
      "         0.      ,   1.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   1.      ,\n",
      "         1.      ,   1.      ,   1.      ,   1.      ,   0.      ,\n",
      "         0.      ,   1.      ,   1.      ,   1.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   1.      ,\n",
      "         0.      ,   1.      ,   0.      ,   0.      ,   0.      ,\n",
      "         1.      ,   1.      ,   0.      ,   0.      ,   1.      ,\n",
      "         0.      ,   0.      ,  -1.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   1.      ,   1.      ,   1.      ,\n",
      "         1.      ,   1.      ,   0.      ,   0.      ,   1.      ,\n",
      "         1.      ,   1.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   1.      ,   0.      ,   1.      ,\n",
      "         0.      ,   0.      ,   0.      ,   1.      ,   1.      ,\n",
      "         0.      ,   0.      ,   1.      ,   0.      ,   0.      ,\n",
      "         1.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         0.      ,   0.      ,   0.      ,   0.      ,   0.      ,\n",
      "         6.      ,   1.      ,   2.      ,   0.      ], dtype=float32), 'action_mask': array([1., 1., 1., 0., 0., 0., 0., 1., 1., 0., 0., 0., 0., 1., 0., 0., 0.,\n",
      "       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "       0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 1., 0., 0., 0.,\n",
      "       0., 1.], dtype=float32)}, 0, False, {'depth': 1})\n"
     ]
    }
   ],
   "source": [
    "print(env.step(28))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(env.step(42))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.it_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset(file=\"../../Dataset/function031563/function031563_generator.cpp\")\n",
    "env.step(2)\n",
    "env.step(4)\n",
    "env.step(1)\n",
    "env.step(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(env.obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset(\"../../Dataset/function550010/function550010_generator.cpp\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prog=Tiramisu_Program(\"../../Dataset_copies/function550398/function550398_generator.cpp\")\n",
    "annotations=prog.get_program_annotations()\n",
    "rep_temp = get_representation_template(annotations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(100):\n",
    "    print(i)\n",
    "    env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(100):\n",
    "    print(i)\n",
    "    env.reset()\n",
    "    env.step(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(5):\n",
    "    f.write(\"---------- Episode {} ------------------\\n\\n\".format(i))\n",
    "    obs=env.reset(\"../../Dataset_copies/function550205/function550205_generator.cpp\")\n",
    "    done=False\n",
    "    step=0\n",
    "    while not done:\n",
    "        f.write(\"---------- Step {} ------------------\\n\".format(step))\n",
    "        raw_action=random.randint(0,5)\n",
    "        print(\"the raw action is {}\\n\".format(raw_action))\n",
    "        f.write(\"Action {}: \\n\".format(raw_action))\n",
    "        obs, reward, done, info = env.step(raw_action)\n",
    "        \n",
    "        f.write(\"Schedule str: {}\\n\".format(env.schedule_str))\n",
    "        f.write(\"Schedule: {}\\n\".format(env.schedule))\n",
    "        f.write(\"Observation: {}\\n\".format(obs))\n",
    "        f.write(\"Reward: {}\\n\".format(reward))\n",
    "        f.write(\"Done: {}\\n\".format(done))\n",
    "        f.write(\"Info: {}\\n\".format(info))\n",
    "        step+=1\n",
    "       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "compile_and_run_tiramisu_code(\"../Dataset_copies/function551346/function551346_solver.cpp\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "action=Action(0,it_dict_prog)\n",
    "action_params=action.parameter()\n",
    "print(action_params)\n",
    "params=[]\n",
    "print(len(action_params['first_dim_index']))\n",
    "print(action_params['second_dim_index']==2)\n",
    "if len(action_params['first_dim_index'])==2 and len(action_params['second_dim_index'])==2:\n",
    "    params.append(int(action_params['first_dim_index'][-1]))\n",
    "    params.append(int(action_params['second_dim_index'][-1]))\n",
    "print(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Path(\"/home/narimane/Desktop/PFE/RL_in_Search_space_exploration/Dataset_copies/function550010\").parts[1:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = open(\"../../new_dataset.json\")\n",
    "progs_dict=json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = {}\n",
    "d[0]=(0, 1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d[0][2] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "obs=gym.spaces.Box(low=-np.inf, high=np.inf, shape=(5,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs=np.array([0, 1, 0, \"NaN\", 1], dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([False, False, False,  True, False])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.isnan(obs)==True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.11 ('env_pytorch')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "050826732b6dd14ee3ecea8886821d9d0b77488371b776336303352bdd6bb0ce"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
